﻿#pragma once
#include "Configuration.h"
#include "Types.h"
#include <vtkm/cont/ExecutionAndControlObjectBase.h>
#include <vtkm/Math.h>

//#include <vtkm/cont/ExecutionObjectBase.h>

struct EquationContext : vtkm::cont::ExecutionAndControlObjectBase
{
  static constexpr UInt8 NUM_VARIABLES = 5;
  using ValueType = Vec<Real, NUM_VARIABLES>;

  VTKM_CONT EquationContext(const Configuration& cfg)
    : _mach(cfg.Get<Real>("mach", 0.1))
  {
  }

  VTKM_EXEC_CONT ValueType PrimaryToConservation(const ValueType& primary_value) const
  {
    ValueType value;
    auto r = primary_value[0];
    auto u = primary_value[1];
    auto v = primary_value[2];
    auto w = primary_value[3];
    auto p = primary_value[4];

    value[0] = r;
    value[1] = r * u;
    value[2] = r * v;
    value[3] = r * w;
    value[4] = p / _gamm1 + 0.5 * r * (u * u + v * v + w * w);
    return value;
  }

  VTKM_EXEC_CONT ValueType ConservationToPrimary(const ValueType& conservation_value) const
  {
    ValueType value;
    auto r = conservation_value[0];
    auto u = conservation_value[1] / r;
    auto v = conservation_value[2] / r;
    auto w = conservation_value[3] / r;
    Real p = _gamm1 * (conservation_value[4] - r * (u * u + v * v + w * w)/2.);

    return { r, u, v, w, p };
  }

  VTKM_EXEC_CONT
  void InvisTerm(const ValueType& uh, Vec3f* inv_term)
  {
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };
    FloatDefault p = _gamm1 * (uh[4] - 0.5 * uh[0] * vtkm::Dot(vel, vel));
    FloatDefault h = (uh[4] + p) / uh[0];

    inv_term[0] = uh[0] * vel;

    inv_term[1] = uh[1] * vel;
    inv_term[1][0] += p;

    inv_term[2] = uh[2] * vel;
    inv_term[2][1] += p;

    inv_term[3] = uh[3] * vel;
    inv_term[3][2] += p;

    inv_term[4] = uh[0] * h * vel;
  }

  VTKM_EXEC_CONT
  void InvisTermFlux(const ValueType& uh, const Vec3f& normal, ValueType& flux)
  {
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };
    FloatDefault p = _gamm1 * (uh[4] - 0.5 * uh[0] * vtkm::Dot(vel, vel));
    FloatDefault h = (uh[4] + p) / uh[0];

    Real vn = vtkm::Dot(vel, normal);
    flux[0] = uh[0] * vn;
    flux[1] = uh[1] * vn + p * normal[0];
    flux[2] = uh[2] * vn + p * normal[1];
    flux[3] = uh[3] * vn + p * normal[2];
    flux[4] = uh[0] * h * vn;
  }

  VTKM_EXEC_CONT void RiemannFlux(const ValueType& ul,
                                  const ValueType& ur,
                                  const Vec3f& normal,
                                  ValueType& riemann_flux)
  {
    ValueType flux_l, flux_r;
    InvisTermFlux(ul, normal, flux_l);
    InvisTermFlux(ur, normal, flux_r);

    ValueType average = ul; //(ul + ur) / 2.;
    auto perimary_value = ConservationToPrimary(average);
    Vec3f vel{ perimary_value[1], perimary_value[2], perimary_value[3] };
    Real vn = vtkm::Dot(vel, normal);
    Real c = vtkm::Sqrt(_gamma * perimary_value[4] / perimary_value[0]);
    Real lam = (vtkm::Abs(vn) + c);
    //console::Info("lam:", lam);

    riemann_flux = (flux_l + flux_r) / 2. + 0.2*lam * (ul - ur);

    this->ul = ul;
    this->ur = ur;
  }


  EquationContext PrepareForExecution(vtkm::cont::DeviceAdapterId, vtkm::cont::Token&) const
  {
    return *this;
  }

  EquationContext PrepareForControl() const { return *this; }

  Real _mach;
  Real _gamma = 1.4, _gamm1 = _gamma - 1;
  ValueType ul, ur;
  ValueType fluxl, fluxr;
};
